import os
import logging
import time
from datetime import timedelta
from avalanche.evaluation import Metric
from typing import List, Union
import shutil


class LogFormatter(logging.Formatter):

    def __init__(self):
        super(LogFormatter, self).__init__()
        self.start_time = time.time()

    def format(self, record):
        elapsed_seconds = round(record.created - self.start_time)

        prefix = "%s - %s - %s" % (
            record.levelname,
            time.strftime("%x %X"),
            timedelta(seconds=elapsed_seconds),
        )
        message = record.getMessage()
        message = message.replace(
            "\n", "\n" + " " * (len(prefix) + 3)
        )
        return "%s - %s" % (prefix, message)


def create_logger(output_dir, name, vb=2):

    logger = LOGGER(output_dir, name, vb)

    return logger


class LOGGER:
    def __init__(self, output_dir, name, vb=2):
        """
        Create a logger.
        """
        self.output_dir = output_dir
        self.name = name

        # create log formatter
        self.log_formatter = LogFormatter()

        # create file handler and set level to debug
        log_file = os.path.join(output_dir, name + ".log")
        self.file_handler = logging.FileHandler(log_file, "w")
        self.file_handler.setLevel(logging.DEBUG)
        self.file_handler.setFormatter(self.log_formatter)

        # create console handler and set level to info
        self.console_handler = logging.StreamHandler()
        log_level = (
            logging.DEBUG
            if vb == 2
            else logging.INFO if vb == 1 else logging.WARNING
        )
        self.console_handler.setLevel(log_level)
        self.console_handler.setFormatter(self.log_formatter)

        # create logger and set level to debug
        self.logger = logging.getLogger()
        self.logger.handlers = []
        self.logger.setLevel(logging.DEBUG)
        self.logger.propagate = False
        self.logger.addHandler(self.file_handler)
        self.logger.addHandler(self.console_handler)

    def reset_time(self):
        self.log_formatter.start_time = time.time()

    def info(self, message):
        self.logger.info(message)

    def debug(self, message):
        self.logger.debug(message)

    def warning(self, message):
        self.logger.warning(message)

    def backup(self, bp_name):
        # copy old log file
        self.file_handler.close()
        self.console_handler.close()

        new_log_file = os.path.join(self.output_dir, bp_name + ".log")
        log_file = os.path.join(self.output_dir, self.name + ".log")
        shutil.copy(log_file, new_log_file)

    def format_info(self, dicts):
        for k, v in dicts.items():
            self.logger.info(f"{k:<15s}: {v}")


class cusMetric(Metric):
    """Custom Metric

    Args:
        Metric (Metric): Metric
    """

    def __init__(self, name: str = "") -> None:
        self.name: str = name
        self.summed: float = 0.0
        self.weight: float = 0.0

    def result(self):
        if self.weight == 0.0:
            return 0.0
        return self.summed / self.weight

    def reset(self):
        """
        Resets the metric.

        :return: None.
        """
        self.summed = 0.0
        self.weight = 0.0

    def update(self, value, weight=1.0):
        value = float(value)
        weight = float(weight)
        self.summed += value * weight
        self.weight += weight

    def __add__(self, other: "cusMetric") -> "cusMetric":
        """
        Return a metric representing the weighted mean of the 2 means.

        :param other: the other mean
        :return: The weighted mean"""
        res = cusMetric()
        res.summed = self.summed + other.summed
        res.weight = self.weight + other.weight
        return res

    @property
    def value(self) -> float:
        if self.weight == 0.0:
            return 0.0
        return self.summed / self.weight


class Metrics:
    """Format multi Metrics

    Attributes:
        metrics (List[cusMetric]): list of metrics
        format_str (str): format string
    """

    def __init__(self, metrics: List[cusMetric], prefix="  ") -> None:
        self.prefix = prefix
        self.metrics = metrics

    def update(
        self, values: List, weights: Union[List[float], int] = 1
    ) -> None:
        weights_list = (
            [weights] * len(values)
            if isinstance(weights, int)
            else weights
        )
        for m, v, w in zip(self.metrics, values, weights_list):
            m.update(v, w)

    def reset(self) -> None:
        for m in self.metrics:
            m.reset()

    def display(self) -> str:
        print(self.format_str)
        return self.format_str

    @property
    def format_str(self) -> str:
        return " ".join(
            [
                self.prefix,
                *[
                    (
                        " " * len(self.prefix)
                        + f"{i.name}: {i.result():.3e}\n"
                    )
                    for i in self.metrics
                ],
            ]
        )
